"""
基于 LangChain 的投资决策智能体
"""
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from agents.base_agent import StockAgent
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Literal
from colorama import Fore, Style
from enum import Enum
import json


class InvestmentRecommendation(str, Enum):
    """投资建议枚举"""
    STRONG_BUY = "强烈买入"
    BUY = "买入"
    HOLD = "持有"
    SELL = "卖出"
    STRONG_SELL = "强烈卖出"


class InvestmentDecision(BaseModel):
    """投资决策结果模型"""
    recommendation: InvestmentRecommendation = Field(description="投资建议等级")
    confidence_score: float = Field(description="置信度评分 (0-100)", ge=0, le=100)
    target_price: str = Field(description="目标价格区间")
    time_horizon: str = Field(description="投资时间周期")
    key_reasons: List[str] = Field(description="主要决策理由")
    risk_factors: List[str] = Field(description="主要风险因素")
    action_plan: str = Field(description="具体操作建议")
    summary: str = Field(description="决策总结")


class DecisionMakerAgent(StockAgent):
    """基于 LangChain 的投资决策智能体"""
    
    def __init__(self):
        super().__init__("decision_maker")
        self.decision_chain = self._create_decision_chain()
        self.output_parser = PydanticOutputParser(pydantic_object=InvestmentDecision)

    def _create_system_prompt(self) -> str:
        """重写系统提示词"""
        return """你是一个资深的投资决策专家，具有多年的股票投资和风险管理经验。

你的决策职责包括：
1. 综合分析技术面和基本面信息
2. 评估投资风险和收益潜力
3. 制定明确的投资建议和操作策略
4. 提供合理的目标价格和时间周期
5. 识别关键的风险控制点

决策原则：
- 基于数据和事实进行客观判断
- 充分考虑风险因素
- 提供可执行的操作建议
- 保持谨慎和理性的投资态度

投资建议等级说明：
- 强烈买入：预期收益>20%，风险可控
- 买入：预期收益10-20%，风险适中
- 持有：预期收益0-10%，维持现状
- 卖出：预期收益<0%，建议减仓
- 强烈卖出：预期大幅下跌，建议清仓"""

    def _create_decision_chain(self):
        """创建决策处理链"""
        
        # 创建决策提示词模板
        decision_prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template(self._create_system_prompt()),
            HumanMessagePromptTemplate.from_template("""
请基于以下分析报告做出投资决策：

{analysis_report}

请提供以下内容：
1. 明确的投资建议等级（强烈买入/买入/持有/卖出/强烈卖出）
2. 置信度评分（0-100分）
3. 目标价格区间
4. 建议的投资时间周期
5. 主要决策理由（至少3个要点）
6. 主要风险因素（至少2个要点）
7. 具体的操作建议
8. 决策总结

请确保决策逻辑清晰，建议具体可行。
""")
        ])
        
        def enhance_decision(decision_result: str) -> str:
            """增强决策结果"""
            enhanced_template = """
# 投资决策报告

{decision_result}

---
**决策框架说明：**
- 使用 LangChain 智能决策链
- 基于多维度风险收益分析
- 结合技术面和基本面判断
- 采用量化评分体系

**重要提醒：**
1. 本决策建议基于当前可获得的信息
2. 市场存在不确定性，请结合个人风险承受能力
3. 建议分批建仓，设置止损点
4. 定期回顾和调整投资策略

**免责声明：**
本报告仅供参考，不构成具体的投资建议。投资者应当根据自身情况独立决策，并承担相应风险。
"""
            return enhanced_template.format(decision_result=decision_result)
        
        # 创建处理链
        chain = (
            {"analysis_report": RunnablePassthrough()}
            | decision_prompt
            | self.llm
            | StrOutputParser()
            | RunnableLambda(enhance_decision)
        )
        
        return chain

    def _create_structured_decision_chain(self):
        """创建结构化决策链"""
        
        structured_prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template(
                self._create_system_prompt() + "\n\n" + 
                "请按照以下JSON格式返回决策结果：\n{format_instructions}"
            ),
            HumanMessagePromptTemplate.from_template("""
基于以下分析报告做出投资决策：
{analysis_report}
""")
        ])
        
        chain = (
            {
                "analysis_report": RunnablePassthrough(),
                "format_instructions": lambda _: self.output_parser.get_format_instructions()
            }
            | structured_prompt
            | self.llm
            | self.output_parser
        )
        
        return chain

    def execute(self, analysis_report: str) -> str:
        """执行投资决策任务"""
        try:
            print(Fore.MAGENTA + f"\n    - [LangChain] 启动投资决策链...")
            
            # 使用 LangChain 决策链处理
            decision_result = self.decision_chain.invoke(analysis_report)
            
            print(Fore.GREEN + "    - [LangChain] 投资决策链处理完成")
            return decision_result
            
        except Exception as e:
            error_msg = f"投资决策过程中发生错误: {str(e)}"
            print(Fore.RED + f"    - [Error] {error_msg}")
            return error_msg

    def execute_structured_decision(self, analysis_report: str) -> Dict[str, Any]:
        """执行结构化投资决策（返回结构化结果）"""
        try:
            print(Fore.MAGENTA + f"\n    - [LangChain] 启动结构化决策链...")
            
            structured_chain = self._create_structured_decision_chain()
            result = structured_chain.invoke(analysis_report)
            
            print(Fore.GREEN + "    - [LangChain] 结构化决策完成")
            return result.dict()
            
        except Exception as e:
            error_msg = f"结构化决策过程中发生错误: {str(e)}"
            print(Fore.RED + f"    - [Error] {error_msg}")
            return {"error": error_msg}

    def get_decision_summary(self, analysis_report: str) -> Dict[str, Any]:
        """获取决策摘要"""
        try:
            structured_result = self.execute_structured_decision(analysis_report)
            if "error" in structured_result:
                return structured_result
            
            return {
                "recommendation": structured_result["recommendation"],
                "confidence": structured_result["confidence_score"],
                "summary": structured_result["summary"]
            }
        except Exception as e:
            return {"error": f"获取决策摘要失败: {str(e)}"}